package cn.yx.framework.security;

import javax.crypto.*;
import javax.crypto.spec.DESKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;

public class Encorypter {
    public static final String ALGORITHM_DES = "DES";
    public static final String ALGORITHM_DESEDE = "DESede";
    public static final String ALGORITHM_AES = "AES";
    public static final String ALGORITHM_BLOWFISH = "Blowfish";
    public static final String ALGORITHM_RC2 = "RC2";
    public static final String ALGORITHM_RC4 = "RC4";

    public Encorypter() {
    }

    public byte[] encrypt(String algorithm, byte[] key, byte[] data) throws InvalidKeySpecException, InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, BadPaddingException, IllegalBlockSizeException {
        Key k = this.getKey(algorithm, key);
        Cipher cipher = Cipher.getInstance(algorithm);
        cipher.init(1, k);
        return cipher.doFinal(data);
    }

    public byte[] decrypt(String algorithm, byte[] key, byte[] data) throws InvalidKeyException, InvalidKeySpecException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
        Key k = this.getKey(algorithm, key);
        Cipher cipher = Cipher.getInstance(algorithm);
        cipher.init(2, k);
        return cipher.doFinal(data);
    }

    public String encrypt(String algorithm, String source, String key) throws NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, InvalidKeyException, InvalidKeySpecException {
        return new String(this.encrypt(algorithm, source.getBytes(), key.getBytes()));
    }

    public byte[] initKey(String algorithm) throws NoSuchAlgorithmException {
        return this.initKey(algorithm, (byte[])null);
    }

    public byte[] initKey(String algorithm, byte[] seed) throws NoSuchAlgorithmException {
        SecureRandom secureRandom;
        if(seed != null) {
            secureRandom = new SecureRandom(seed);
        } else {
            secureRandom = new SecureRandom();
        }

        KeyGenerator kg = KeyGenerator.getInstance(algorithm);
        kg.init(secureRandom);
        return kg.generateKey().getEncoded();
    }

    private Key getKey(String algorithm, byte[] key) throws InvalidKeyException, InvalidKeySpecException, NoSuchAlgorithmException {
        if("DES".equals(algorithm)) {
            DESKeySpec dks = new DESKeySpec(key);
            SecretKeyFactory keyFactory = SecretKeyFactory.getInstance(algorithm);
            return keyFactory.generateSecret(dks);
        } else {
            return new SecretKeySpec(key, algorithm);
        }
    }
}
